"""
RequestId 中间件
为每个请求添加或获取 RequestId
"""
import uuid
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response


class RequestIdMiddleware(BaseHTTPMiddleware):
    """
    RequestId 中间件
    
    如果请求头中没有 X-Request-Id，则自动生成一个
    并在响应头中返回 X-Request-Id
    """
    
    async def dispatch(self, request: Request, call_next):
        # 获取或生成 RequestId
        request_id = request.headers.get("X-Request-Id")
        if not request_id:
            request_id = str(uuid.uuid4())
        
        # 将 RequestId 存储到请求状态中，方便后续使用
        request.state.request_id = request_id
        
        # 调用下一个中间件或路由处理
        response: Response = await call_next(request)
        
        # 在响应头中添加 RequestId
        response.headers["X-Request-Id"] = request_id
        
        return response

